from ..config import OpPrior, Space, autoSimply

__all__ = [
    "Operable",
    "OpBase",
    "Op1",
    "Op2",
    "OpF1",
    "OpF2",
    "OpConst",
]

_autoSimp = lambda sth: sth.simple() if autoSimply else sth


class Desc:
    def __init__(self, index):
        self.index = index

    def __get__(self, obj: "OpBase", _):
        return obj.args[self.index]

    def __set__(self, obj: "OpBase", val):
        obj.args[self.index] = val


class Operable:
    def __init__(self, val=None):
        self.value = val

    def __bool__(self):
        return True

    def eval(self, **val):
        return self.getValue(val)

    def getValue(self, valDict: dict):
        return self.value

    def diff(self, dx) -> "Operable":
        return OpConst(0)

    def simple(self) -> "Operable":
        return self

    def forward(self):
        return self

    def backward(self) -> None:
        pass

    @classmethod
    def same(cls, *arg):
        return cls(*arg)


class OpConst(Operable):
    def __init__(self, val):
        assert val is not None
        super().__init__(val)

    def __str__(self):
        if isinstance(self.value, float):
            return f"{self.value:.3f}"
        else:
            return str(self.value)

    def __bool__(self):
        return bool(self.value)
    
    def __eq__(self, other) -> bool:
        return other == self.value or super().__eq__(other)


class OpBase(Operable):
    op = ''
    disp = ''
    lamda = lambda *_: None
    argsNum = 0
    level = OpPrior.default
    addSpace = Space.default

    def __init__(self, *args):
        assert len(args) == self.argsNum
        self.args = [i if isinstance(i, Operable) else OpConst(i) for i in args]
        super().__init__(
            self.lamda(*(i.value for i in self.args))
            if all(i.value is not None for i in self.args)
            else None
        )

    def __iter__(self):
        return self.__next__()

    def __next__(self):
        for arg in self.args:
            if isinstance(arg, OpBase):
                yield from arg
        yield self

    def getValue(self, valDict: dict):
        return self.lamda(*(i.getValue(valDict) for i in self.args))

    def forward(self):
        for arg in self.args:
            arg.forward()
            assert arg.value is not None
        self.value = self.lamda(*(i.value for i in self.args))
        return self

    def backward(self):
        pass

    @classmethod
    def strWithBrackets(cls, obj) -> str:
        return (
            f"({obj})"
            if isinstance(obj, OpBase) and obj.level < cls.level
            else str(obj)
        )


class Op1(OpBase):
    argsNum = 1
    r = Desc(0)

    def __str__(self):
        r = self.strWithBrackets(self.r)
        return f"{self.disp} {r}" if self.addSpace else f"{self.disp}{r}"

    def simple(self):
        r = self.r.simple()
        if isinstance(r, OpConst):
            return OpConst(self.lamda(r.value))
        return self.same(r)

    @classmethod
    def getOP(cls):
        return lambda r: _autoSimp(cls(r))


class Op2(OpBase):
    argsNum = 2
    l, r = Desc(0), Desc(1)

    def __str__(self):
        l = self.strWithBrackets(self.l)
        r = self.strWithBrackets(self.r)
        return f"{l} {self.disp} {r}" if self.addSpace else f"{l}{self.disp}{r}"

    def simple(self):
        l, r = self.l.simple(), self.r.simple()
        if isinstance(l, OpConst) and isinstance(r, OpConst):
            return OpConst(self.lamda(l.value, r.value))
        return self.same(l, r)

    @classmethod
    def getOPs(cls):
        def OPL(l, r):
            return _autoSimp(cls(l, r))

        def OPR(r, l):
            return _autoSimp(cls(l, r))

        return OPL, OPR


class OpF1(Op1):
    level = OpPrior.func
    addSpace = Space.func

    @staticmethod
    def strWithBrackets(obj) -> str:
        return f"({obj})"


class OpF2(Op2):
    level = OpPrior.func
    addSpace = Space.func

    def __str__(self):
        return f"{self.disp}({self.l}," + (
            f" {self.r})" if self.addSpace else f"{self.r})"
        )
